#!/usr/bin/env python
import numpy as np
import itk
import scipy.ndimage as ndi
import imageTools.ITKUtils.io as io
import sys
import os

def lung_extraction(fname): 
    reader = itk.ImageFileReader.IF3.New()
    reader.SetFileName(fname)
    reader.Update()
    btf = itk.BinaryThresholdImageFilter.IF3IUS3.New()
    btf.SetInput(reader.GetOutput())
    btf.SetLowerThreshold(-900)
    btf.SetUpperThreshold(-200)
    ccf = itk.ConnectedComponentImageFilter.IUS3IUL3.New()
    ccf.SetInput(btf.GetOutput())
    ccf.SetFullyConnected(1)
    rcf = itk.RelabelComponentImageFilter.IUL3IUS3.New()
    rcf.SetInput(ccf.GetOutput())
    btf2 = itk.BinaryThresholdImageFilter.IUS3IUS3.New()
    btf2.SetInput(rcf.GetOutput())
    btf2.SetLowerThreshold(1)
    btf2.SetUpperThreshold(1)
    btf2.Update()
    writer = itk.ImageFileWriter.IUS3.New()
    tmp = os.path.splitext(fname)
    fout = tmp[0] + "_lung" + tmp[1] 
    writer.SetFileName(fout) 
    writer.SetInput(btf2.GetOutput())
    writer.Update()
    return fout

def applyMask(fname, mask):
    mask_img = io.readImage(mask)
    array = itk.PyBuffer.IF3.GetArrayFromImage(mask_img)
    mask = np.where( array > 0, 1,0).astype(np.uint8)
    img = io.readImage(fname) 
    img_array_1 = itk.PyBuffer.IF3.GetArrayFromImage(img)
    img_array = mask*img_array_1
    tmp = os.path.splitext(fname)
    outputfile = tmp[0] + "_masked" + tmp[1]  
    output_img = itk.PyBuffer.IF3.GetImageFromArray(img_array)
    io.writeImage(output_img, outputfile) 
    return outputfile
